package download

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/grafov/m3u8"
	"github.com/orestonce/gopool"
	"github.com/spf13/cobra"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"ts-downloader/entity"
	"ts-downloader/log"
)

var (
	m3u8Url, saveDir string
	threadCount      int
	Cmd              = &cobra.Command{
		Use: "download",
		Run: run,
	}
)

func init() {
	Cmd.Flags().StringVarP(&m3u8Url, "m3u8Url", "u", "", "M3u8Url")
	Cmd.Flags().StringVarP(&saveDir, "saveDir", "d", "", "文件保存路径(默认为当前路径)")
	Cmd.Flags().IntVarP(&threadCount, "threadCount", "", 8, "下载线程数")
}

func run(cmd *cobra.Command, args []string) {
	if m3u8Url == "" {
		_ = cmd.Help()
		return
	}
	d := NewDownloader()
	d.download(cmd.Context())
}

type Downloader struct {
	logFileLocker sync.Mutex
}

func NewDownloader() *Downloader {
	d := &Downloader{}
	return d
}

func (d *Downloader) download(ctx context.Context) {
	u, err := url.Parse(m3u8Url)
	checkErr(err)
	err = d.downloadExec(ctx, u)
	if err != nil {
		log.Error("err = %v", err)
		return
	}
	fmt.Println("下载成功")
}

func (d *Downloader) downloadExec(ctx context.Context, m3u8Url *url.URL) error {
	playList, err := sniffM3u8(ctx, m3u8Url.String())
	if err != nil {
		return err
	}
	log.Info("playList len = %d", len(playList.Segments))

	var encInfo *entity.EncryptInfo
	getTsInfoFunc := func(ts *m3u8.MediaSegment) (entity.TsInfo, error) {
		var tsInfo = entity.TsInfo{
			MediaSegment: ts,
			EncryptInfo:  encInfo,
		}
		if ts.Key != nil {
			encInfo, err = getEncryptInfo(ctx, ts)
			if err != nil {
				return tsInfo, err
			}
			log.Info("key changed %+v, key = %s", ts.Key.URI, string(encInfo.Key))
		}
		if tsUrl, err := url.Parse(tsInfo.URI); err != nil {
			log.Error("ts.Url parse err %v", err)
			return tsInfo, err
		} else {
			if tsUrl.Scheme == "" {
				tsUrl = m3u8Url.ResolveReference(tsUrl)
				tsInfo.URI = tsUrl.String()
				log.Info("ts url  %v", tsInfo.URI)
			}
		}
		return tsInfo, nil
	}

	downloadingDir, err := d.initDir(saveDir)
	if err != nil {
		return err
	}
	taskPool := gopool.NewThreadPool(threadCount)
	asyncExec := func(tsInfo entity.TsInfo) {
		outputTsFilePath := fmt.Sprintf("%s/%05d.ts", downloadingDir, tsInfo.SeqId)
		taskPool.AddJob(func() {
			if err = d.downloadTsFile(ctx, tsInfo, outputTsFilePath); err != nil {
				log.Error("downloadTsFile err = %v", err)
			}
		})
	}

	for _, ts := range playList.Segments {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		if ts == nil {
			break
		}
		tsInfo, err := getTsInfoFunc(ts)
		if err != nil {
			return err
		}
		asyncExec(tsInfo)
	}
	taskPool.CloseAndWait()
	return nil
}

// 下载ts文件
// @modify: 2020-08-13 修复ts格式SyncByte合并不能播放问题
func (d *Downloader) downloadTsFile(ctx context.Context, ts entity.TsInfo, outputTsFilePath string) (err error) {
	var stat os.FileInfo
	stat, err = os.Stat(outputTsFilePath)
	if err == nil && stat.Mode().IsRegular() {
		log.Info("ts [%s] file is exist", ts.URI)
		return nil
	}
	res, err := DoGetRequest(ctx, ts.URI, 10)
	if err != nil {
		log.Error("[warn] download fail :%s", ts.URI)
		return err
	}
	// 校验长度是否合法
	origData := res.Bytes()
	// 解密出视频 ts 源文件
	if ts.EncryptInfo != nil {
		//解密 ts 文件，算法：aes 128 cbc pack5
		origData, err = ts.Decode(ts.SeqId, origData)
		if err != nil {
			log.Error("AesDecrypt fail, ts = %v", ts.URI)
			return err
		}
	}
	// Detect Fake png file
	if bytes.HasPrefix(origData, entity.PngSign) {
		origData = origData[8:]
	}

	// https://en.wikipedia.org/wiki/MPEG_transport_stream
	// Some TS files do not start with SyncByte 0x47, they can not be played after merging,
	// Need to remove the bytes before the SyncByte 0x47(71).
	//syncByte := uint8(71) //0x47
	//bLen := len(origData)
	//for j := 0; j < bLen; j++ {
	//	if origData[j] == syncByte {
	//		origData = origData[j:]
	//		break
	//	}
	//}
	tmpPath := outputTsFilePath + ".tmp"
	err = os.WriteFile(tmpPath, origData, 0666)
	if err != nil {
		return err
	}
	err = os.Rename(tmpPath, outputTsFilePath)
	return err
}

func checkErr(e error) {
	if e != nil {
		log.Panic(e)
	}
}

func (d *Downloader) initDir(saveDir string) (string, error) {
	tmp1 := sha256.Sum256([]byte(m3u8Url))
	downloadingDir := filepath.Join(saveDir, "downloading", hex.EncodeToString(tmp1[:]))
	if !isDirExists(downloadingDir) {
		err := os.MkdirAll(downloadingDir, os.ModePerm)
		if err != nil {
			log.Error("os.MkdirAll error0: %s ", err.Error())
			return "", err
		}
	}
	return downloadingDir, nil
}

func sniffM3u8(ctx context.Context, u8Url string) (res *m3u8.MediaPlaylist, err error) {
	m3u8Url := u8Url
	response, err := DoGetRequest(ctx, m3u8Url, 3)
	if err != nil {
		log.Error("[warn] File :%s", m3u8Url)
		return nil, err
	}
	if !response.Ok {
		return nil, errors.New(strconv.Itoa(response.StatusCode))
	}
	payList, listType, err := entity.M3u8Parse(response.Bytes())
	if err != nil {
		log.Error("[warn] M3u8Parse :%s", m3u8Url)
		return nil, err
	}
	switch listType {
	case m3u8.MEDIA:
		mediapl := payList.(*m3u8.MediaPlaylist)
		return mediapl, nil
	case m3u8.MASTER:
		//_ = payList.(*m3u8.MasterPlaylist)
		return nil, errors.New("请访问m3u8子资源")
	default:
		return nil, errors.New("unknown err")
	}
}

// 获取m3u8加密的密钥, 可能存在iv
func getEncryptInfo(ctx context.Context, segment *m3u8.MediaSegment) (info *entity.EncryptInfo, err error) {
	uri := segment.Key.URI
	if uri == "" {
		return nil, nil
	}
	method := segment.Key.Method
	if method == entity.EncryptMethod_NONE {
		return nil, nil
	}
	keyUrl, err := resolveRefUrl(m3u8Url, uri)
	if err != nil {
		return nil, err
	}
	res, err := DoGetRequest(ctx, keyUrl, 3)
	if err != nil {
		log.Error("[warn] getEncryptInfo url = %s, err  = %v", keyUrl, err)
		return nil, err
	}
	if !res.Ok {
		return nil, errors.New(res.String())
	}

	if method == entity.EncryptMethod_AES128 && len(res.Bytes()) != 16 { // Aes 128
		return nil, errors.New("getEncryptInfo invalid key " + strconv.Quote(res.String()))
	}

	return entity.NewEncryptInfo(method, segment.Key.IV, res.Bytes())
}

func isDirExists(path string) bool {
	stat, err := os.Stat(path)
	return err == nil && stat.IsDir()
}
func resolveRefUrl(baseUrl string, extUrl string) (after string, err error) {
	urlObj, err := url.Parse(baseUrl)
	if err != nil {
		return "", err
	}
	lineObj, err := url.Parse(extUrl)
	if err != nil {
		return "", err
	}
	return urlObj.ResolveReference(lineObj).String(), nil
}
